#include <kontiki/trajectory_estimator.h>
#include <kontiki/trajectories/split_trajectory.h>
#include <kontiki/trajectories/uniform_r3_spline_trajectory.h>
// #include <kontiki/trajectories/uniform_so3_spline_trajectory.h>
#include <kontiki/measurements/orientation_measurement.h>
#include <kontiki/measurements/position_measurement.h>
#include <kontiki/measurements/velocity_measurement.h>

#include <Eigen/Core>
#include <fstream>
#include <memory>
#include <vector>
#include <ceres/ceres.h>

#include <ros/ros.h>
#include <tf/transform_datatypes.h>
#include <tf/transform_broadcaster.h>
#include <nav_msgs/Path.h>
#include <nav_msgs/Odometry.h>

#include "encode_spline/TicToc.h"


using R3TrajEstimator         = kontiki::TrajectoryEstimator<kontiki::trajectories::UniformR3SplineTrajectory>;
using SplitTrajEstimator = kontiki::TrajectoryEstimator<kontiki::trajectories::SplitTrajectory>;
using PositionMeasurement     = kontiki::measurements::PositionMeasurement;
using OrientationMeasurement  = kontiki::measurements::OrientationMeasurement;
using VelocityMeasurement     = kontiki::measurements::VelocityMeasurement;

using Result = std::unique_ptr<kontiki::trajectories::TrajectoryEvaluation<double>>;
using namespace std;
std::shared_ptr<kontiki::trajectories::SplitTrajectory> traj_;
std::shared_ptr<R3TrajEstimator> estimator_r3;  // 用于优化
std::shared_ptr<SplitTrajEstimator> estimator_split;  // 用于优化
// std::shared_ptr<PositionMeasurement> m_p0_; // 初始化t0处的位置测量，为（0 0 0）
const double PI = 3.1415926;


void initialTrajTo(double max_time) {
    Eigen::Vector3d p0(0,0,0);
    Eigen::Quaterniond q0 = Eigen::Quaterniond::Identity();
    traj_->R3Spline()->ExtendTo (max_time, p0);
    traj_->SO3Spline()->ExtendTo(max_time, q0);
    std::cout << "spline min time = " << traj_->R3Spline()->MinTime() << "  max time = " << traj_->R3Spline()->MaxTime() << std::endl;
    std::cout << "spline contral point num = " << traj_->R3Spline()->NumKnots() << std::endl;
}

bool evaluate(double t, Eigen::Vector3d &p) {
    if (traj_->MinTime() > t || traj_->MaxTime() <= t)
        return false;
    Result result = traj_->Evaluate(t, kontiki::trajectories::EvalPosition);
    p = result->position;
    return true;
}

bool evaluateRot(double t, Eigen::Quaterniond &q) {
    if (traj_->MinTime() > t || traj_->MaxTime() <= t)
        return false;
    Result result = traj_->Evaluate(t, kontiki::trajectories::EvalOrientation);
    q = result->orientation;
    return true;
}

inline void pubOdometry(ros::Publisher& publisher,
                        const Eigen::Vector3d &p, 
                        const Eigen::Quaterniond &q,
                        const ros::Time &timestamp, 
                        nav_msgs::Odometry &local_odom_) {

  local_odom_.header.stamp = timestamp;
  local_odom_.header.seq += 1;
  local_odom_.pose.pose.orientation.x = q.x();
  local_odom_.pose.pose.orientation.y = q.y();
  local_odom_.pose.pose.orientation.z = q.z();
  local_odom_.pose.pose.orientation.w = q.w();
  local_odom_.pose.pose.position.x = p.x();
  local_odom_.pose.pose.position.y = p.y();
  local_odom_.pose.pose.position.z = p.z();
  publisher.publish(local_odom_);

}

template<typename T>
bool IsUnitQuaternion(const Eigen::Quaternion<T>& q) {
  auto err = ceres::abs(q.norm() - T(1));
  return err < T(1e-5);
}

void test_encode() {
    // 测试数据，1st是时间，2nd是角度值
    vector<pair<double, double>> data;
    double traj_end_time = 2;
    double knot_distance = 0.05;
    double traj_start_time = 0;
    // double velo = 0.0;
    // double a = 300.0 / 180.0 * PI;
    // int count = traj_end_time / 0.01;
    // double last = 0; 
    // for (int i = 0; i < count + 1; ++i) {
    //     double value = last + velo * 0.01 + 0.5 * a * 0.01 * 0.01;
    //     double ra = (rand() % 5 + 1) / 1000;
    //     ra = 0.;
    //     double time = 0.01 * i + ra;
    //     data.push_back({time, value});
    //     last = value;
    //     velo = velo + a * 0.01;
    // }

    // 读取实际数据
    ifstream infile;
    infile.open("/home/dut-zxw/zxw/lidar_encode_calib_ws/measure.txt");
    double data_t, data_v;
    while (infile >> data_t >> data_v) {
        data.push_back({data_t, data_v});
    }
    infile.close();

    traj_start_time = data[0].first;
    traj_end_time = data.back().first + 0.02;
    traj_ = std::make_shared<kontiki::trajectories::SplitTrajectory>
            (knot_distance, knot_distance, traj_start_time, traj_start_time);
    initialTrajTo(traj_end_time);    

    // 添加测量值
    for (int i = 0; i < 3; ++i) {
        estimator_r3 = std::make_shared<R3TrajEstimator>(traj_->R3Spline());
        // data[0].first += 0.0001;
        std::vector< std::shared_ptr<PositionMeasurement>>  msg_list_; // 保存测量值，需要保存智能指针
        for (auto da : data) {
            double time = da.first;
            double value = da.second;
            Eigen::Vector3d p0(value, 0, 0);
            double weight = 1.0;
            auto m_p0_= std::make_shared<PositionMeasurement>(time, p0, weight);
            msg_list_.push_back(m_p0_);
            estimator_r3->template AddMeasurement<PositionMeasurement>(m_p0_);
        }    

        // 求解
        ceres::Solver::Summary summary = estimator_r3->Solve(50, false);
        std::cout << summary.BriefReport() << std::endl;
    }

    // 输出样条曲线上的值
    vector<pair<double, double>> spline_value;
    for (double i = traj_start_time; i < traj_end_time; i += 0.01) {
        double time = i;
        Eigen::Vector3d p;
        if (evaluate(time, p)) {
            spline_value.push_back({time, p.x()});
        } else {
            cout << "time error\n";
        }
    }
    cout << "write file, data.size = " << data.size() << "  spline_value.size = " << spline_value.size() << endl;;
    ofstream outfile;
    // 测量数据
    outfile.open("/home/dut-zxw/zxw/test_any/test_spline_ws/measure.txt");
    for (auto v : data) {
        outfile << v.first << " " << v.second << endl;
    }
    outfile.close();
    // 样条估计数据
    outfile.open("/home/dut-zxw/zxw/test_any/test_spline_ws/spline.txt");

    for (auto v : spline_value) {
        outfile << v.first << " " << v.second << endl;
    }

    return;
}

int main(int argc, char **argv) {
    ros::init(argc, argv, "spline");

    ros::NodeHandle n("~");

    ros::Time time; 
    time = ros::Time::now();

	ros::Publisher path_pub = n.advertise<nav_msgs::Path>("/trajectory",1, true);
	ros::Publisher Bpath_pub = n.advertise<nav_msgs::Path>("/spline",1, true);

	ros::Publisher pub_measure = n.advertise<nav_msgs::Odometry>("/measure", 100);
    ros::Publisher pub_opt = n.advertise<nav_msgs::Odometry>("/optimization", 100);
	ros::Publisher pub_control = n.advertise<nav_msgs::Odometry>("/control", 100);

    nav_msgs::Odometry measure_odom, opt_odom, control_odom;
	measure_odom.header.frame_id = "world";
  	measure_odom.child_frame_id = "lidar";
	opt_odom.header.frame_id = "world";
  	opt_odom.child_frame_id = "lidar";
    control_odom.header.frame_id = "world";
  	control_odom.child_frame_id = "lidar";

    vector<Eigen::Vector3d> postions;
    vector<Eigen::Quaterniond> rotatinos;
    vector<double> times;
    const double pi = 3.1415926;
    int one_circle = 360.0 / 3.0;
    double delta_theta = 2.0 * pi / one_circle;
    int nums = pi / delta_theta / 2.0;
    double delta_time = 0.1;
    double knot_distance = 1;
    postions.resize(nums);
    rotatinos.resize(nums);
    times.resize(nums);
    double r = 3;
    double theta = pi / 4.0, t = 1;
    Eigen::Vector3d mp;
    Eigen::Quaterniond mq;
    Eigen::Vector3d zaxis(0, 0, 1);
    for (int i = 0; i < nums; ++i) {

        // 获得测量值
        mp.x() = r * sin(theta);
        mp.y() = r * cos(theta);
		mp.z() = 0;

        mq = Eigen::Quaterniond(Eigen::AngleAxisd(-theta, zaxis).toRotationMatrix());

        t += delta_time;

        postions[i] = mp;
        rotatinos[i] = mq;
        times[i] = t;

        theta += delta_theta;
        r += 0.03;
    }

    double traj_start_time = times[0];
    double traj_end_time = times.back();
    traj_ = std::make_shared<kontiki::trajectories::SplitTrajectory>
            (knot_distance, knot_distance, traj_start_time, traj_start_time);
    initialTrajTo(traj_end_time);  

    estimator_split = std::make_shared<SplitTrajEstimator>(traj_);
    std::vector< std::shared_ptr<PositionMeasurement>>  msg_list_; // 保存测量值，需要保存智能指针
    std::vector< std::shared_ptr<OrientationMeasurement>> orien_list_;
    
    for (int i = 0; i < times.size(); ++i) {
        double time = times[i];
        double weight = 1.0;
        auto m_p0_= std::make_shared<PositionMeasurement>(time, postions[i], weight);
        msg_list_.push_back(m_p0_);
        estimator_split->template AddMeasurement<PositionMeasurement>(m_p0_);  

        auto mg_rot = std::make_shared<OrientationMeasurement>(time, rotatinos[i], weight);
        orien_list_.push_back(mg_rot);
        estimator_split->template AddMeasurement<OrientationMeasurement>(mg_rot);
    }

    // 求解
    TicToc tictoc;
    tictoc.Tic();
    ceres::Solver::Summary summary = estimator_split->Solve(50, false);
    std::cout << summary.BriefReport() << std::endl;
    cout << "solve time : " << tictoc.Toc() << " ms\n";
    cout << "measurement size : " << times.size() << endl;
    cout << "knot nums : " << traj_->R3Spline()->NumKnots() << endl;

    vector<Eigen::Vector3d> opt_pos;
    vector<Eigen::Quaterniond> opt_rot;
    vector<double> opt_time;
    delta_time /= 2.0;
    for (double i = times[0]; i < times.back(); i += delta_time) {
        evaluate(i, mp);
        opt_pos.push_back(mp);
        evaluateRot(i, mq);
        opt_rot.push_back(mq);
        opt_time.push_back(i);
        if (!IsUnitQuaternion(mq)) {
            cout << "not a unit: " << mq.coeffs().transpose() << "  pos: " << mp.transpose() << " time: " << i << endl;
        }
    }

    cout << "get opt measurements\n";
    
    ros::Rate rate(10);
    for (int i = 0, j = 0; i < times.size(); ++i) {
        mp = postions[i];
        mq = rotatinos[i];
        ros::Time rt = ros::Time::now();
        pubOdometry(pub_measure, mp, mq, rt, measure_odom);

        for (; j < opt_time.size() && opt_time[j] < times[i]; ++j) {
            mp = opt_pos[j];
            mq = opt_rot[j];
            rt = ros::Time::now();
            pubOdometry(pub_opt, mp, mq, rt, opt_odom);
        }
        rate.sleep();
    }
    
    vector<Eigen::Vector3d> control_pos;
    vector<Eigen::Quaterniond> control_rot;
    for (int i = 0; i < traj_->R3Spline()->NumKnots(); ++i) {
        mp = traj_->R3Spline()->GetContralPointaAt(i);
        mq = traj_->SO3Spline()->GetContralPointaAt(i);
        control_pos.push_back(mp);
        control_rot.push_back(mq);
        ros::Time rt = ros::Time::now();
        pubOdometry(pub_control, mp, mq, rt, control_odom);
        Eigen::Quaterniond pidentity = rotatinos[5*i];
        traj_->SO3Spline()->SetContralPointAt(i, pidentity);
        rate.sleep();
    }

    // 以下为测试只优化四元数
    
    orien_list_.clear();
    estimator_split = std::make_shared<SplitTrajEstimator>(traj_);
    for (int i = 0; i < times.size(); ++i) {
        double time = times[i];
        double weight = 1.0;

        auto mg_rot = std::make_shared<OrientationMeasurement>(time, rotatinos[i], weight);
        orien_list_.push_back(mg_rot);
        estimator_split->template AddMeasurement<OrientationMeasurement>(mg_rot);
    }
    summary = estimator_split->Solve(50, false);
    std::cout << summary.BriefReport() << std::endl;
    cout << "solve time : " << tictoc.Toc() << " ms\n";
    cout << "measurement size : " << times.size() << endl;
    cout << "knot nums : " << traj_->R3Spline()->NumKnots() << endl;

    opt_pos.clear();
    opt_rot.clear();
    opt_time.clear();
    delta_time /= 2.0;
    for (double i = times[0]; i < times.back(); i += delta_time) {
        evaluate(i, mp);
        opt_pos.push_back(mp);
        evaluateRot(i, mq);
        opt_rot.push_back(mq);
        opt_time.push_back(i);
        if (!IsUnitQuaternion(mq)) {
            cout << "not a unit: " << mq.coeffs().transpose() << "  pos: " << mp.transpose() << " time: " << i << endl;
        }
    }

    cout << "get opt measurements\n";
    for (int i = 0; i < opt_time.size(); ++i) {
        mp = opt_pos[i];
        mq = opt_rot[i];
        ros::Time rt = ros::Time::now();
        pubOdometry(pub_opt, mp, mq, rt, opt_odom);
        rate.sleep();
    }

    cout << "print enter to end\n";
    getchar();

    cout << "==============================\nend\n";
    return 0;
}

